/*
MIT License Copyright (c) [2024] [Chuck Fu]

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/
#pragma once
#ifndef CF_CRYPTO_SM2_H
#define CF_CRYPTO_SM2_H
#include "stdio.h"
#include "stdlib.h"
#include "stdint.h"
#include "stdbool.h"

// -----------------------------------------------------------------------------
// Utility Struct - 256Bit BigInteger
// A self-contained implementation of 256 bit bigint with the operation allowed
// in the prime Field Fp.
// -----------------------------------------------------------------------------
typedef enum bigint_cmp_enum {
    BIGINT_CMP_LESS   = -1,
    BIGINT_CMP_EQUAL   = 0,
    BIGINT_CMP_GREATER = 1,
} bigint_cmp_order;
// Defines the bigint type, the detail is hidden from the user.
typedef struct uint256_st bigint;
// create bigint object
bigint* bigint_new();
bigint* bigint_new_from_hex_str(const char *str);
bigint* bigint_new_from_u32(uint32_t x);
bigint* bigint_new_from_u64(uint64_t x);
bigint* bigint_new_from_other(const bigint* other);
bigint* bigint_new_from_bytes(const uint8_t* bytes);
// delete bigint object
void bigint_delete(bigint* x);
// set bigint values
void bigint_set_from_u32(bigint* x, uint32_t val);
void bigint_set_from_other(bigint* x, const bigint* other);
// check if its zero
bool bigint_is_zero(const bigint* x);
// get the byte numbers needed to represent the bigint.
int bigint_get_byte_len(const bigint *x);
// compare a and b, the result is -1 (less), 0(same), 1(greater)
bigint_cmp_order bigint_cmp(const bigint* a, const bigint* b);
// convert the bigint to byte array.
void bigint_to_bytes(bigint* x, uint8_t *bytes);

// -----------------------------------------------------------------------------
// Bigint Math Functions in Fp:
// It supports: Add, Sub, Mul, Inv, Pow operations.
// Note: Div can be implemented with the inverse mod, e.g. a / b = a * Inv(b)
// -----------------------------------------------------------------------------
// t = t % p, this normalized t to the range [0, p)
void bigint_norm(bigint* t, const bigint* p);
// t = (a + b) % p
void bigint_add_mod_p(
    bigint* t,
    const bigint* a,
    const bigint* b,
    const bigint* p
);
// t = (a - b) % p
void bigint_sub_mod_p(
    bigint* t,
    const bigint* a,
    const bigint* b,
    const bigint* p
);
// t = (a * b) % p
void bigint_mul_mod_p(
    bigint* t,
    const bigint* a,
    const bigint* b,
    const bigint* p
);
// t = (a^b) % p
void bigint_pow_mod_p(
    bigint* t,
    const bigint* a,
    const bigint* b,
    const bigint* p
);
// t = a^-1 % p
void bigint_inv_mod_p(
    bigint*,
    const bigint* a,
    const bigint* p
);

// -----------------------------------------------------------------------------
// Utility Struct - Bit String
// The bitstring struct stores bits in byte array, and keep track of the active
// bit counts.
// -----------------------------------------------------------------------------
typedef struct bitstring_st bitstring;
// create bitstring object
bitstring* bitstring_new(uint64_t bytes);
bitstring* bitstring_new_from_cstr(const char* cstr);
bitstring* bitstring_new_from_hex_str(const char *str);
bitstring* bitstring_new_from_bytes(const uint8_t* bytes, int byte_len);
// delete bitstring object
void bitstring_delete(bitstring*);
// get the byte used to store all the bits.
uint64_t bitstring_get_byte_len(const bitstring* bs);
// get the active bit count
uint64_t bitstring_get_bit_len(const bitstring* bs);
// convert to bigint (in SM2 ctx, this must be 256 bits)
bigint* bitstring_to_bigint(const bitstring* bs);

// -----------------------------------------------------------------------------
// SM3 hash
// -----------------------------------------------------------------------------
/**
 * Hash function used in SM2, map any length bit stream to 256 bits hash value.
 *
 * @param s represents the bit stream
 * @return the hash value
 */
bitstring* sm3_hash256(const bitstring* s);

// -----------------------------------------------------------------------------
// SM2 Interfaces
// -----------------------------------------------------------------------------
/**
 * struct sm2_ctx_st SM2 Running Context
 *
 * SM2 use the ECC (Elliptic Curve Crypto) curve in the 256 bit prime field
 * with the following function: y^2 = x^3 + ax + b
 *
 * @p: prime number, represent the field Fp
 * @a: the parameter a in the ECC equation
 * @b: the parameter b in the ECC equation
 * @n: the order of base point G
 * @Gx: x cord of base point G
 * @Gy: y cord of base point G
 * @Px: x cord of public key point P
 * @Py: x cord of public key point P
 * @d: private key used
 */
typedef struct sm2_ctx_st {
    bigint* p;
    bigint* a;
    bigint* b;
    bigint* n;
    // Base Point
    bigint* Gx;
    bigint* Gy;
    // Public Key
    bigint* Px;
    bigint* Py;
    // Private Key
    bigint* d;
} sm2_ctx;

/**
 * struct sm_signature_st
 * @r: the r value in the signature
 * @s: the s value in the signature
 */
typedef struct sm2_signature_st {
    bitstring* r;
    bitstring* s;
} sm2_signature;

/**
 * Encode SM2 Signature to DER (Distinguished Encoding Rules)
 * The ASN.1 def of SM2 signature is
 *
 * SM2Signature :: = SEQUENCE {
 *     r INTEGER,
 *     s INTEGER
 * }
 *
 * @param signature the SM2 signature <r, s> pair
 * @return The encoded DER types stored in @bitstring
 */
bitstring* sm2_signature_to_der(const sm2_signature* signature);

// SM2 functions
/**
 * Create an empty SM2 context object.
 *
 * @return the created SM2 ctx object
 */
sm2_ctx* sm2_ctx_new();
/**
 * Init the SM2 context with default parameters values.
 *
 * @param ctx the ctx object used
 */
void sm2_ctx_init(sm2_ctx* ctx);

/**
 * Delete the SM2 context object.
 *
 * @param ctx the ctx object used
 */
void sm2_ctx_delete(sm2_ctx* ctx);

/**
 * Set the private key used in the context.
 *
 * @param ctx the ctx object used
 * @param private_key the private_key to set
 */
void sm2_ctx_set_private_key(sm2_ctx* ctx, bigint* private_key);

/**
 * Set the public key used in the context.
 *
 * @param ctx the ctx object used
 * @param x the x cord of public key point to set
 * @param y the y cord of public key point to set
 */
void sm2_ctx_set_public_key(sm2_ctx* ctx, bigint* x, bigint* y);


// SM2 ECC Solver
/**
 * struct sm_ecc_point - Represent a vector/point in ECC space.
 * Note: When used in Affine Space, z cord represent the infinity point when
 *       z is not zero.
 * The implementation detail is hidden from user, only use
 * sm2_ecc_point_new_from_xy(x, y) to create a new ECC point in user space.
 *
 * @x the x cord of point
 * @y the y cord of point
 * @z the z cord of point
 */
typedef struct sm2_ecc_point_st {
    bigint *x, *y, *z;
} sm2_ecc_point;

/**
 * Create a 0 point in ECC space (the infinity point)
 *
 * @return the created point in ECC space
 */
sm2_ecc_point* sm2_ecc_point_new();

/**
 * Create a point in ECC space use (x,y).
 *
 * @param x the x cord of the point
 * @param y the y cord of the point
 * @return the created point in ECC space
 */
sm2_ecc_point* sm2_ecc_point_new_from_xy(bigint* x, bigint* y);

/**
 * Delete an ECC point.
 *
 * @param p the point to delete
 */
void sm2_ecc_point_delete(sm2_ecc_point* p);

/**
 * Point ADD operation in ECC space.
 *
 * @param ctx the ctx object used
 * @param res the result of the add operation, res = P1 + P2
 * @param p1 the point used in the add operation
 * @param p2 the point used in the add operation
 */
void sm2_ecc_point_add(
    const sm2_ctx* ctx,
    const sm2_ecc_point* res,
    const sm2_ecc_point* p1,
    const sm2_ecc_point* p2
);

/**
 * Point MUL k operation in ECC space.
 *
 * @param ctx the ctx object used
 * @param p the point used in the add operation
 * @param k the point used in the add operation
 *
 * @return the result of the sum of k p's in ECC space. res = kP
 */
sm2_ecc_point* sm2_ecc_point_mul_k(
    const sm2_ctx* ctx,
    const sm2_ecc_point* p,
    const bigint* k
);

/**
 * Create the user distinguishing identifier
 * ZA=H256(ENTLA ∥ IDA ∥ a ∥ b ∥ Gx ∥ Gy ∥ Px ∥ Py)
 *
 * @param ctx the ctx object used
 * @param id_a the user provided string
 * @return the user distinguishing identifier
 */
bitstring* sm2_new_za(const sm2_ctx* ctx, const bitstring* id_a);

/**
 * Random bigint integer k generation, used in sign process.
 *
 * @param ctx the ctx object used
 * @return the random biginteger used in sign process
 */
bigint* sm2_new_k(const sm2_ctx* ctx);

/**
 * Create the digest of za and message.
 *
 * @param za the user distinguishing identifier
 * @param message the message to sign
 * @return za = H256(za||message)
 */
bitstring* sm2_digest(
    const bitstring* za,
    const bitstring* message
);

// SM2 Sign
/**
 * struct sm2_sign_result is used to encapsulate the result of an SM2 signature
 * operation.
 *
 * when @pass is true, @signature is the result
 * when @pass is false, @error_message stores the error message
 */
typedef struct sm2_sign_result_st {
    bool pass;
    char* error_message;
    sm2_signature* signature;
} sm2_sign_result;

/**
 * Generate an SM2 signature against a given message and the user distinguishing
 * identifier.
 *
 * @param ctx the ctx object used
 * @param za the user distinguishing identifier
 * @param message the message received
 * @return the signature generated
 */
sm2_signature* sm2_sign(
    const sm2_ctx* ctx,
    const bitstring* za,
    const bitstring* message
);

/**
 * Generate an SM2 signature against a fix k value and digest
 *
 * @param ctx the ctx object used
 * @param digest the result of H256(za||message)
 * @param digest the random number k used in the sign process
 * @return the signature generated
 */
sm2_sign_result* sm2_sign_with_k(
    const sm2_ctx* ctx,
    const bitstring* digest,
    const bigint* k
);

/**
 * Delete the sign result object.
 */
void sm2_sign_result_delete(sm2_sign_result*);

/**
 * struct sm2_verify_result is used to encapsulate the result of an SM2 verify
 * operation.
 *
 * when @pass is false, @error_message stores the error message
 * when @pass is true, means the signature pass the verification.
 */
typedef struct sm2_verify_result_st {
    bool pass;
    char* error_message;
} sm2_verify_result;

/**
 * Verifies an SM2 signature against a given message and the user distinguishing
 * identifier.
 *
 * @param ctx the ctx object used
 * @param digest the result of H256(za||message)
 * @param signature the signature to verify
 * @return the verified result
 */
sm2_verify_result* sm2_verify(
    const sm2_ctx* ctx,
    const bitstring* digest,
    const sm2_signature* signature
);

/**
 * Delete the verify result object.
 */
void sm2_verify_result_delete(sm2_verify_result*);

// -----------------------------------------------------------------------------
// Debug interfaces
// -----------------------------------------------------------------------------
// Print to std output
void show_bitstring(const bitstring* bs);
void show_bigint(const bigint* b);
void show_bigint_as_uint64(const bigint *x);
void show_sm2_ctx(const sm2_ctx* ctx);
void show_sm2_ecc_point(const char* point_name, const sm2_ecc_point* p);
void show_sm2_signature(const sm2_signature* signature);
// Write to files
void write_bitstring(FILE* file, const bitstring* bs);
void write_bigint(FILE* file, const bigint* x);
#endif //CF_CRYPTO_SM2_H
